import json
import time
from typing import Callable
import requests
from enum import Enum
import threading

class RequestType(Enum):
    STARTING = 0
    COMPLETING = 1
    COMPLETED = 2

class DeepseekClient:
    def __init__(self, api_key: str):
        self.key = api_key

        self.char_queue = []
        self.completed = False

    def chatComplete(self, prompt: str, callback: Callable[[RequestType, str], None]) -> None:
        self.completed = False
        def _chatComplete(prompt: str, callback: Callable[[RequestType, str], None]):
            callback(RequestType.STARTING, '')
            response = requests.post(
                url='https://api.deepseek.com/chat/completions',
                headers={
                    'Content-Type': 'application/json',
                    'Accept': 'application/json',
                    'Authorization': f'Bearer {self.key}'
                },
                data=json.dumps({
                    'messages': [
                        {
                            'role': 'system',
                            'content': f'You are a helpful assistant. Please respond to the user very briefly.(because width of the screen is limited). And please don\'t use markdown. Now time is {time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())}. Use English to reply.'
                        },
                        {
                            'role': 'user',
                            'content': prompt
                        }
                    ],
                    'model': 'deepseek-chat',
                    'stream': True
                }),
                stream=True
            )

            def _callback():
                while True:
                    if self.char_queue:
                        callback(RequestType.COMPLETING, self.char_queue.pop(0))
                        if len(self.char_queue) > 0:
                            time.sleep(0.6 / len(self.char_queue))
                    if self.completed and not self.char_queue:
                        callback(RequestType.COMPLETED, '')
                        break

            back_thread = threading.Thread(target=_callback)
            back_thread.start()

            for line in response.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    line = line[6:]
                    if line == '[DONE]':
                        self.completed = True
                        break
                    else:
                        data = json.loads(line)
                        content = data['choices'][0]['delta']['content']
                        for char in content:
                            self.char_queue.append(char)

        thread = threading.Thread(target=_chatComplete, args=(prompt, callback))
        thread.start()